#!/usr/bin/env python3

import logging
import sys
import os
import argparse
import subprocess

# 优化：将R脚本路径定义为常量，方便修改和管理
SEURAT_TO_LOOM_R_SCRIPT = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/py/SingleCell/utils4cmd/pysceinc/seurat2loom_exprMat.r"
SCEINC_RES_TO_RDS_R_SCRIPT = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/py/SingleCell/utils4cmd/pysceinc/sceinc_res_loom2rds.r"

# 移除 conda 激活函数
# 优化：不再尝试在脚本中激活conda环境，建议用户手动激活或使用 conda run
def activate_conda_env(env_name):
    # 构建激活环境的命令
    command = f"source /Pub/Users/wangyk/miniconda3/bin/activate {env_name} && env"

    # 使用 subprocess.run 执行命令并捕获输出
    result = subprocess.run(command, capture_output=True,
                            text=True, shell=True, executable="/bin/bash")

    if result.returncode != 0:
        print(f"Error activating conda environment: {result.stderr}")
        return

    # 解析环境变量
    env_vars = result.stdout.splitlines()
    for var in env_vars:
        key, value = var.split("=", 1)
        os.environ[key] = value

activate_conda_env("/Pub/Users/wangyk/miniconda3/envs/pyscenic")


wrp = '''
利用pyscenic进行转录因子识别，只需要seurat对象路径以及结果输出路径。
默认为人类数据集。

小鼠数据集：
   如果使用小鼠单细胞数据集，请务必替换以下三个文件参数 (-tfs, -f, -tbl) 
   为小鼠对应的参考文件。
   参考资源: https://github.com/aertslab/pySCENIC/tree/master/resources
            https://resources.aertslab.org/cistarget/databases/mus_musculus/mm9/refseq_r39/mc9nr/gene_based/

使用参考:
   建议先手动激活 pyscenic conda 环境，再运行此脚本。
   例如:
   conda activate pyscenic
   python pySCENIC_script.py -s seurat_obj.rds -od sceinc_out/

参数说明:
'''

parser = argparse.ArgumentParser(prog='pySCENIC',
                                 description=wrp,
                                 formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument('-s', '--seurat_obj', help="Seurat对象路径，rds或者rdata",
                    default="/Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/out/SC_part/05.macro_sub_cellchat/ss_m_sub_anno_ready.rds", metavar='')
parser.add_argument('-l', '--loom', help="单细胞loom文件, 和seurat对象一样，用于sceinc分析，二者有一个就行，优先使用loom文件。",
                    default=None, metavar='')
parser.add_argument('-od', '--output_dir', help="结果输出路径",
                    default='/Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/test/', metavar='')
parser.add_argument('-t', '--type',
                    help="分析所用数据类型，貌似seurat中的counts或者data都可以。默认data",
                    default='data',
                    metavar='')
parser.add_argument('-tfs', '--tfs',
                    help="转录因子列表文件路径 (txt)。\n  人类默认值 (hs_hgnc_tfs.txt), 参考: https://github.com/aertslab/pySCENIC/tree/master/resources",
                    default="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/hs_hgnc_tfs.txt",
                    metavar='')
parser.add_argument('-f', '--feather',
                    help="feather 文件路径。\n  人类默认值 (hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather), 参考: https://resources.aertslab.org/cistarget/databases/homo_sapiens/hg19/refseq_r45/mc9nr/gene_based/hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather",
                    default="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather",
                    metavar='')
parser.add_argument('-tbl', '--tbl',
                    help="Motif-TF 对应关系文件路径 (tbl)。\n  人类默认值 (motifs-v9-nr.hgnc-m0.001-o0.0.tbl), 参考: https://resources.aertslab.org/cistarget/motif2tf/motifs-v9-nr.hgnc-m0.001-o0.0.tbl",
                    default="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/motifs-v9-nr.hgnc-m0.001-o0.0.tbl",
                    metavar='')
parser.add_argument('-o', '--organism',
                    help="物种，默认值 'hsa' (人类)。如果是鼠类数据，请设置为 'mmu'，并确保-tfs, -f, -tbl 参数已替换为鼠类文件",
                    default="hsa",
                    metavar='')
# 优化：添加 num_workers 参数，允许用户自定义 pyscenic 的 workers 数量
parser.add_argument('-nw', '--num_workers',
                    help="pyscenic 并行worker数量，默认为 24",
                    default="24",
                    metavar='')


args = parser.parse_args()

tfs = args.tfs
feather = args.feather
tbl = args.tbl
base_files = {tfs, feather, tbl}
od = args.output_dir
input_loom = f"{args.output_dir}/seurat_obj_expression.loom"
num_workers = args.num_workers # 从参数获取 workers 数量

def get_logger(name: str):
    """
    输入logger名称即可创建logger对象。
    """
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)

    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)

    return logger

logger = get_logger(name='pySCENIC')


def ensure_directory_exists(directory_path: str):
    """检查并创建输出目录."""
    if not os.path.exists(directory_path):
        logger.info(f"创建输出目录: {directory_path}") # 优化：添加日志信息
        os.makedirs(directory_path, exist_ok=True) # 优化：使用 exist_ok=True，避免并发问题


def file_check(file_path: str):
    """检查文件是否存在，不存在则报错并退出."""
    if not os.path.exists(file_path):
        logger.error(f"文件不存在: {file_path}")
        sys.exit(1)
    else:
        logger.info(f"文件检查通过: {file_path}") # 优化：添加日志信息


def run_subprocess_command(command_list, step_name):
    """
    运行 subprocess 命令并进行错误处理和日志记录。

    Args:
        command_list (list): 完整的命令列表，包括程序和参数。
        step_name (str):  步骤名称，用于日志记录。

    Returns:
        subprocess.CompletedProcess: subprocess.run 的结果对象。
                                     如果执行失败，则程序会退出。
    """
    logger.info(f"开始执行步骤: {step_name}") # 优化：更清晰的日志信息
    command_str = " ".join(command_list) # 优化：打印命令，方便debug
    logger.info(f"执行命令: {command_str}")

    result = subprocess.run(command_list, capture_output=True, text=True)

    if result.returncode != 0:
        logger.error(f"步骤 {step_name} 执行失败，程序退出.") # 优化：更清晰的错误日志
        logger.error(f"错误信息 (stderr):\n{result.stderr}") # 优化：打印错误输出
        logger.error(f"标准输出 (stdout):\n{result.stdout}") # 优化：打印标准输出，有时错误信息可能在 stdout
        sys.exit(1)
    else:
        logger.info(f"步骤 {step_name} 执行成功.") # 优化：更清晰的成功日志
        if result.stderr: # 优化：即使成功也打印 stderr，以防有警告信息
            logger.warning(f"步骤 {step_name} 警告信息 (stderr):\n{result.stderr}\n")
        return result


def main():
    ensure_directory_exists(args.output_dir)

    logger.info("开始 pySCENIC 分析流程...") # 优化：添加总流程开始日志

    logger.info("检查输入文件...") # 优化：文件检查步骤日志
    for i in base_files:
        file_check(i)

    if args.loom is not None:
        # os.copy_file_range(args.loom, input_loom)
        os.system(f"cp {args.loom} {args.output_dir}seurat_obj_expression.loom")
    else:
        # 1. Seurat to Loom
        logger.info("1. Seurat 对象转换为 loom 文件") # 优化：步骤序号和更友好的步骤名称
        r_seurat2loom_exprMat_path_command = [
            SEURAT_TO_LOOM_R_SCRIPT,
            args.seurat_obj,
            args.output_dir,
            '-t', args.type
        ]
        run_subprocess_command(r_seurat2loom_exprMat_path_command, "Seurat to Loom 转换")


    # 2. pySCENIC grn
    logger.info("2.1. 运行 pySCENIC grn (基因调控网络推断)") # 优化：步骤序号和更友好的步骤名称
    grn_command = [
        "pyscenic", "grn",
        "--num_workers", num_workers, # 优化：使用参数控制 worker 数量
        "--output", f"{od}/adj.sample.tsv",
        "--method", "grnboost2",
        input_loom,
        tfs
    ]
    run_subprocess_command(grn_command, "pySCENIC grn")


    # 3. pySCENIC ctx
    logger.info("2.2. 运行 pySCENIC ctx (转录因子靶基因分析)") # 优化：步骤序号和更友好的步骤名称
    cistarget_command = [
        "pyscenic", "ctx",
        f"{od}/adj.sample.tsv", feather,
        "--annotations_fname", tbl,
        "--expression_mtx_fname", input_loom,
        "--mode", "dask_multiprocessing",
        "--output", f"{od}/reg.csv",
        "--num_workers", num_workers, # 优化：使用参数控制 worker 数量
        "--mask_dropouts"
    ]
    run_subprocess_command(cistarget_command, "pySCENIC ctx")


    # 4. pySCENIC aucell
    logger.info("2.3. 运行 pySCENIC aucell (AUCell活性评分计算)") # 优化：步骤序号和更友好的步骤名称
    aucell_command = [
        "pyscenic", "aucell", input_loom, f"{od}/reg.csv",
        "--output", f"{od}/pySCENIC_res.loom",
        "--num_workers", num_workers # 优化：使用参数控制 worker 数量
    ]
    run_subprocess_command(aucell_command, "pySCENIC aucell")


    # 5. Loom to RDS
    logger.info("2.4. 转换 pySCENIC 结果 loom 文件为 rds 文件") # 优化：步骤序号和更友好的步骤名称
    r_sceinc_res_loom2rds_path_command = [
        SCEINC_RES_TO_RDS_R_SCRIPT,
        f"{od}/pySCENIC_res.loom"
    ]
    run_subprocess_command(r_sceinc_res_loom2rds_path_command, "Loom to RDS 转换")


    logger.info("pySCENIC 分析流程完成!") # 优化：添加总流程完成日志


if __name__ == '__main__':
    main()